#include <memory>
#include <torch/torch.h>
#include <LibDL/data/Example/Example.h>

#ifndef SWIG
typedef std::unique_ptr<
	torch::data::StatelessDataLoader<
	torch::data::datasets::MapDataset<
	torch::data::datasets::MNIST,
	torch::data::transforms::Stack<torch::data::Example<at::Tensor, at::Tensor>>>,
	torch::data::samplers::RandomSampler>,
	std::default_delete<torch::data::StatelessDataLoader<torch::data::datasets::MapDataset<torch::data::datasets::MNIST, 
	torch::data::transforms::Stack<torch::data::Example<at::Tensor, at::Tensor>>>, torch::data::samplers::RandomSampler>>> MNIST_PTR;
#endif

namespace data {
    struct MNISTDataIterator {
        explicit MNISTDataIterator(const MNIST_PTR &mnistPtr) : it(std::move(mnistPtr->begin())),
                                                                end(std::move(mnistPtr->end())) {}

        bool hasNext() const {
            return (it != end);
        }

        const Example nextImpl() {
            Example result = Example(*it);
            ++it;
            return result;
        }

    private:
        torch::data::Iterator<torch::data::Example<at::Tensor, at::Tensor>> it;
        const torch::data::Iterator<torch::data::Example<at::Tensor, at::Tensor>> end;
    };
}
